import os
import glob
import pandas as pd
from pathlib import Path
import time


def excel_to_csv(input_folder, output_folder, sheet_name=0):
    """
    将指定文件夹中的所有 Excel 文件转换为 CSV 格式

    参数:
    input_folder (str): 包含 Excel 文件的文件夹路径
    output_folder (str): CSV 文件的输出文件夹
    sheet_name (str/int): 要转换的工作表名称或索引
    """
    # 确保输出文件夹存在
    Path(output_folder).mkdir(parents=True, exist_ok=True)

    # 获取所有 Excel 文件 (支持 .xls 和 .xlsx)
    excel_files = glob.glob(os.path.join(input_folder, '.[xX][lL][sS]')) + \
                  glob.glob(os.path.join(input_folder, '*.[xX][lL][sS][xX]'))

    if not excel_files:
        print(f"在 {input_folder} 中未找到 Excel 文件")
        return []

    print(f"找到 {len(excel_files)} 个 Excel 文件，开始转换为 CSV...")

    csv_files = []
    for excel_file in excel_files:
        try:
            # 读取 Excel 文件
            df = pd.read_excel(
                excel_file,
                sheet_name=sheet_name,
                engine='openpyxl' if excel_file.lower().endswith('.xlsx') else 'xlrd'
            )

            # 生成 CSV 文件名
            base_name = os.path.splitext(os.path.basename(excel_file))[0]
            csv_file = os.path.join(output_folder, f"{base_name}.csv")

            # 保存为 CSV
            df.to_csv(csv_file, index=False, encoding='utf-8')
            csv_files.append(csv_file)

            print(f"转换成功: {os.path.basename(excel_file)} -> {os.path.basename(csv_file)}")

        except Exception as e:
            print(f"转换 {excel_file} 失败: {str(e)}")

    return csv_files


def merge_csv_by_rows(csv_files, output_file):
    """
    按行合并多个 CSV 文件（纵向合并）

    参数:
    csv_files (list): CSV 文件路径列表
    output_file (str): 合并后的输出文件路径
    """
    if not csv_files:
        print("没有可合并的 CSV 文件")
        return

    print("\n开始按行合并 CSV 文件...")

    # 读取并合并所有 CSV 文件
    dfs = []
    for csv_file in csv_files:
        try:
            df = pd.read_csv(csv_file)
            dfs.append(df)
            print(f"读取: {os.path.basename(csv_file)} ({len(df)} 行, {len(df.columns)} 列)")
        except Exception as e:
            print(f"读取 {csv_file} 失败: {str(e)}")

    if not dfs:
        print("没有成功读取任何 CSV 文件")
        return

    # 纵向合并
    merged_df = pd.concat(dfs, ignore_index=True)

    # 保存合并结果
    merged_df.to_csv(output_file, index=False)

    print(f"\n按行合并完成! 总行数: {len(merged_df)}")
    print(f"结果已保存到: {output_file}")


def merge_csv_by_columns(csv_files, output_file):
    """
    按列合并多个 CSV 文件（横向合并）

    参数:
    csv_files (list): CSV 文件路径列表
    output_file (str): 合并后的输出文件路径
    """
    if not csv_files:
        print("没有可合并的 CSV 文件")
        return

    print("\n开始按列合并 CSV 文件...")

    # 读取所有 CSV 文件
    dfs = []
    for i, csv_file in enumerate(csv_files):
        try:
            df = pd.read_csv(csv_file)
            # 添加源文件名作为前缀到列名
            base_name = os.path.splitext(os.path.basename(csv_file))[0]
            df = df.add_prefix(f"{base_name}_")

            dfs.append(df)
            print(f"读取: {os.path.basename(csv_file)} ({len(df)} 行, {len(df.columns)} 列)")
        except Exception as e:
            print(f"读取 {csv_file} 失败: {str(e)}")

    if not dfs:
        print("没有成功读取任何 CSV 文件")
        return

    # 横向合并
    merged_df = pd.concat(dfs, axis=1)

    # 保存合并结果
    merged_df.to_csv(output_file, index=False)

    print(f"\n按列合并完成! 总列数: {len(merged_df.columns)}")
    print(f"结果已保存到: {output_file}")


if __name__ == "__main__":
    # 配置参数
    INPUT_FOLDER = "C:/Users/xingwenzheng/Desktop/样例数据"  # Excel 文件所在文件夹
    CSV_FOLDER = "C:/Users/xingwenzheng/Desktop/样例数据"  # 转换后的 CSV 存放文件夹
    OUTPUT_FILE_ROWS = "merged_by_rows.csv"  # 按行合并输出文件
    # OUTPUT_FILE_COLS = "merged_by_cols.csv"  # 按列合并输出文件

    # 步骤 1: 将 Excel 转换为 CSV
    start_time = time.time()
    csv_files = excel_to_csv(INPUT_FOLDER, CSV_FOLDER)
    excel_convert_time = time.time() - start_time

    if not csv_files:
        print("没有可处理的文件，程序退出")
        exit()

    # 步骤 2A: 按行合并 (纵向合并 - 适用于结构相同的表格)
    start_time = time.time()
    merge_csv_by_rows(csv_files, OUTPUT_FILE_ROWS)
    row_merge_time = time.time() - start_time

    # 步骤 2B: 按列合并 (横向合并 - 适用于不同文件代表不同列的情况)
    # start_time = time.time()
    # merge_csv_by_columns(csv_files, OUTPUT_FILE_COLS)
    # col_merge_time = time.time() - start_time

    # 打印性能统计
    print("\n性能统计:")
    print(f"Excel 转换耗时: {excel_convert_time:.2f} 秒")
    print(f"按行合并耗时: {row_merge_time:.2f} 秒")
    # print(f"按列合并耗时: {col_merge_time:.2f} 秒")
    print(f"处理文件总数: {len(csv_files)}")